package cn.bugstack.chatgpt.data.domain.weixin.service.Handler.Impl;

import cn.bugstack.chatglm.model.ChatCompletionRequest;
import cn.bugstack.chatglm.model.ChatCompletionSyncResponse;
import cn.bugstack.chatglm.model.Model;
import cn.bugstack.chatglm.model.Role;
import cn.bugstack.chatglm.session.OpenAiSession;
import cn.bugstack.chatgpt.data.domain.weixin.model.entity.MessageTextEntity;
import cn.bugstack.chatgpt.data.domain.weixin.model.entity.UserBehaviorMessageEntity;
import cn.bugstack.chatgpt.data.domain.weixin.service.Handler.MessageHandler;
import cn.bugstack.chatgpt.data.types.sdk.weixin.XmlUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Service;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

@Service
@Slf4j
public class ChatGLMMessageHandler implements MessageHandler {

    @Value("${wx.config.originalid}")
    private String originalId;

    @Resource
    private ThreadPoolExecutor taskExecutor;

    @Resource
    private OpenAiSession openAiSession;

    private final Logger logger = LoggerFactory.getLogger(ChatGLMMessageHandler.class);

    // 存放OpenAi返回结果数据
    private final Map<String, String> openAiDataMap = new ConcurrentHashMap<>();
    // 存放OpenAi调用次数数据
    private final Map<String, Integer> openAiRetryCountMap = new ConcurrentHashMap<>();

    @Override
    public String handleMessage(UserBehaviorMessageEntity message) {
        try {
            logger.info("请求次数：{}", null == openAiRetryCountMap.get(message.getContent().trim()) ? 1 : openAiRetryCountMap.get(message.getContent().trim()));

            // 异步任务【加入超时重试，对于小体量的调用反馈，可以在重试有效次数内返回结果】
            if (openAiDataMap.get(message.getContent().trim()) == null || "NULL".equals(openAiDataMap.get(message.getContent().trim()))) {
                String data = "大模型正在思考中，请再回复我一句【" + message.getContent().trim() + "】";
                // 休眠等待
                Integer retryCount = openAiRetryCountMap.get(message.getContent().trim());
                if (null == retryCount) {// 还没有重试过，那么证明还没有发送信息到大模型，直接发送信息。
                    if (openAiDataMap.get(message.getContent().trim()) == null) {
                        doChatGPTTask(message.getContent().trim());// 调用之后openAidataMap中会存入消息
                    }
                    logger.info("超时重试：{}", 1);
                    openAiRetryCountMap.put(message.getContent().trim(), 1);// 设置这个消息的重试次数是1
                    TimeUnit.SECONDS.sleep(5);//等待五秒，给系统反应的时间，增加鲁棒性
                    new CountDownLatch(1).await();
                } else if (retryCount < 2) {
                    retryCount = retryCount + 1;
                    logger.info("超时重试：{}", retryCount);
                    openAiRetryCountMap.put(message.getContent().trim(), retryCount);
                    TimeUnit.SECONDS.sleep(5);
                    new CountDownLatch(1).await();
                } else {
                    retryCount = retryCount + 1;
                    logger.info("超时重试：{}", retryCount);
                    openAiRetryCountMap.put(message.getContent().trim(), retryCount);
                    TimeUnit.SECONDS.sleep(3);
                    if (openAiDataMap.get(message.getContent().trim()) != null && !"NULL".equals(openAiDataMap.get(message.getContent().trim()))) {
                        data = openAiDataMap.get(message.getContent().trim());
                    }
                }

                // 反馈信息[文本]
                MessageTextEntity res = new MessageTextEntity();
                res.setToUserName(message.getOpenId());
                res.setFromUserName(originalId);
                res.setCreateTime(String.valueOf(System.currentTimeMillis() / 1000L));
                res.setMsgType("text");
                res.setContent(data);

                return XmlUtil.beanToXml(res);
            }

            // 反馈信息[文本]
            MessageTextEntity res = new MessageTextEntity();
            res.setToUserName(message.getOpenId());
            res.setFromUserName(originalId);
            res.setCreateTime(String.valueOf(System.currentTimeMillis() / 1000L));
            res.setMsgType("text");
            res.setContent(openAiDataMap.get(message.getContent().trim()));
            String result = XmlUtil.beanToXml(res);
            logger.info("接收微信公众号信息请求{}完成 {}", message.getOpenId(), result);
            openAiDataMap.remove(message.getContent().trim());
            return result;
        } catch (Exception e) {
            logger.error("接收微信公众号信息请求{}失败 {}", message.getOpenId(), message, e);
            return "";
        }
    }

    public void doChatGPTTask(String content) {

        openAiDataMap.put(content, "NULL");
        //线程池异步处理调用大模型API逻辑
        taskExecutor.execute(() -> {
            // 入参；模型、请求信息；记得更新最新版 ChatGLM-SDK-Java
            ChatCompletionRequest request = new ChatCompletionRequest();
            request.setModel(Model.GLM_3_5_TURBO);
            request.setPrompt(new ArrayList<ChatCompletionRequest.Prompt>() {
                private static final long serialVersionUID = -7988151926241837899L;

                {
                    add(ChatCompletionRequest.Prompt.builder()
                            .role(Role.user.getCode())
                            .content(content)
                            .build());
                }
            });
            // 同步获取结果
            try {
                // 2.1 提供的官网方法
                ChatCompletionSyncResponse response = openAiSession.completionsSync(request);
                // 保存数据
                openAiDataMap.put(content, response.getChoices().get(0).getMessage().getContent());
            } catch (Exception e) {
                throw new RuntimeException(e);
            }

        });
    }
}
